# from torch import nn
# import torch
#
# class JJw(nn.Module):
#     def __init__(self):
#         super().__init__()
#
#     def forward(seld,input):
#         output=input+1
#         return output
#
# JJw=JJw()
# x=torch.tensor(1.0)
# output=JJw(x)
# print(output)
# import torch.nn.functional as F
# import torch
# input=torch.tensor([[1,2,0,3,1],
#                     [0,1,2,3,1],
#                     [1,2,1,0,0],
#                     [5,2,3,1,1],
#                     [2,1,0,1,1]])
# kernel=torch.tensor([[1,2,1],[0,1,0],[2,1,0]])
# input=torch.reshape(input,(1,1,5,5))
# kernel=torch.reshape(kernel,(1,1,3,3))
#
# print(input.shape)
# print(kernel.shape)
#
# output1=F.conv2d(input,kernel,stride=1)
# print(output1)
# output2=F.conv2d(input,kernel,stride=2)
# print(output2)
# output3=F.conv2d(input,kernel,stride=1,padding=1)
# print(output3)
# output4=F.conv2d(input,kernel,stride=2,padding=1)
# print(output4)

import torch
import torchvision
from torch import nn
from torch.nn import Conv2d
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter

dataset=torchvision.datasets.CIFAR10("./dataset",train=False,transform=torchvision.transforms.ToTensor(),
                                     download=True)

dataloader=DataLoader(dataset,batch_size=64)

class JJw(nn.Module):
    def __init__(self):
        super(JJw,self).__init__()
        self.conv1=Conv2d(3,6,3,stride=1,padding=0)

    def forward(self,x):
        x=self.conv1(x)
        return x

writer=SummaryWriter("./logs")
step=0
jjw=JJw()
for data in dataloader:
    imgs,targets=data
    output=jjw(imgs)
    # torch.Size([64, 6, 30, 30])
    writer.add_images("input",imgs,step)
    # torch.Size([64, 3, 32, 32])
    output=torch.reshape(output,shape=(-1,3,30,30))
    writer.add_images("output",output,step)
    step+=1

writer.close()

